from api.constants import TENANT_ID
from api.db import LLMType
from api.db.services.llm_service import LLMBundle


def dup_remove(pre_content_list, now_content_list):
    # 判断元素是否应该被保留
    def should_keep(content, other_list):
        content_prefix = content[:-2]
        return not any(content_prefix in item for item in other_list)

    # 过滤 pre_content_list 和 now_content_list
    filtered_pre = [item for item in pre_content_list if should_keep(item, now_content_list)]
    filtered_now = [item for item in now_content_list if should_keep(item, filtered_pre)]

    return filtered_pre, filtered_now

def getPrompt(text):
    promot = ('你是一个文本切片工具，需要将所给文档内容分段：所给文档内容为：{}'
              '\n[思考内容]\n'
              '1.理解并分析该文档内容，问答类的内容一个问题和其回答切成一段。\n'
              '2.若是含有不同层级的标题，按照最小标题层级的内容切分\n'
              '3.按照内容的不同表达将其分成不同段落，例如文章的正文和（作者、时间、版权信息）等是不同的内容\n'
              '\n[回答要求]\n'
              '1.每个段落应能独立表达完整思想\n'
              '2.每段应包含完整的流程或信息单元，如完整的回答或问题及其解决方案\n'
              '3.将正文内容与目录等内容分开\n'
              '4.不同段落的内容之间用“------”隔开'
              '\n[限制]\n'
              '1.禁止修改和丢失原文本中的任何内容\n'
              '2.禁止输出繁体字\n'
              '3.禁止输出无关内容\n'
              ).format(text)
    return promot


# 对超长内容进行粗略切片
def split_content(content, page_length=4500, tmp_length=2300):
    content_list = []
    start_page = 0
    end_page = page_length
    length_content = len(content)
    for _ in content:
        content_list.append(content[start_page:end_page])
        start_page += tmp_length
        end_page += tmp_length
        if end_page > length_content > start_page:
            content_list.append(content[start_page:-1])
            return content_list
        elif end_page > length_content:
            break
    return content_list


# 对文本进行智能切片
def smart_chunk(question, task):
    content_list = split_content(question)
    print("粗略切片后文本内容:\n", content_list, "\n")
    try:
        llm_name = task["llm_id"]
        tenant_id = task["tenant_id"]
        chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, llm_name)
    except Exception as e:
        print("smart trunk bundle chat model failed", e)
        raise

    result_list = []
    now_cntent_list = []
    pre_content_list = []
    for index, text in enumerate(content_list):
        try:
            msg = [{'role': 'user', 'content': text}]
            prompt = getPrompt(text)
            # gen_conf = {'temperature': 0, 'top_p': 0.3, 'presence_penalty': 0.4, 'frequency_penalty': 0.7,
            #             'max_tokens': 5120}
            gen_conf = {'temperature': 0, 'top_p': 0.3, 'max_tokens': 5120}
            answer = chat_mdl.chat(prompt, msg, gen_conf)
            print("大模型分片去重前内容：\n", answer, "\n")
            now_cntent_list = answer.split("------")

            # 进行去重
            if pre_content_list:
                pre_content_list, now_cntent_list = dup_remove(pre_content_list, now_cntent_list)
                result_list += pre_content_list

            pre_content_list = now_cntent_list
        except Exception as e:
            print("smart trunk called chat model failed", e)
            raise
    result_list += now_cntent_list
    print("大模型分片去重后内容为：", result_list)
    return result_list


if __name__ == "__main__":
    task = {"llm_id": "deepseek-chat@DeepSeek", "tenant_id": TENANT_ID}
    question = "2024" + \
               "2024-05" + \
               "我行携手兴业银行，成功落地首笔台资银行企业数字人民币收款业务" + \
               "2024-05" + \
               "我行统一资信管理平台获得软件著作权" + \
               "2024-03" + \
               "我行首获中央国库现金管理商业银行定期存款参与银行团成员资格" + \
               "2024-01" + \
               "我行首获记账式国债承销团成员资格" + \
               "2024-01" + \
               "我行获批非金融企业债务融资工具承销商资格，进一步丰富我行投行业务品种" + \
               "2023" + \
               "2023-12" + \
               "我行连任中银协台资委主任单位，马立新董事长担任主任委员" + \
               "2023-11" + \
               "我行新核心系统成功投产" + \
               "2023-10" + \
               "我行经过央行数研所、城银清算系统验收获批数字人民币业务资质，并于2023年10月成功上线数字人民币业务" + \
               "2023-08" + \
               "我行获得人民币跨境支付系统（CIPS）直接参与者资格" + \
               "2023-07" + \
               "我行获得债券通「北向通」报价机构资质。可为境外投资者持续提供报价，增加人民币债券参与度，进一步促进人民币国际化" + \
               "2023-07" + \
               "我行获批成为中国外汇交易中心债券通报价机构" + \
               "2023-07" + \
               "我行在全国银行间债券市场成功完成10亿元无固定期限资本债券发行。两期发债总规模共18亿元，为首家在大陆成功发行永续债的台资银行" + \
               "2023-06" + \
               "我行推出“台胞移动支付便利化服务”专案，做好台胞来大陆金融服务第一站" + \
               "2023-05" + \
               "我行济南分行开业" + \
               "2023-05" + \
               "我行受让花旗银行（中国）有限公司个人房产抵押贷款，包括债券及附属担保/保证权益" + \
               "2023-04" + \
               "我行成功完成20亿元普通金融债发行" + \
               "2023-03" + \
               "我行获批成为中国外汇交易中心银行间人民币外汇市场期权尝试做市机构" + \
               "2023-02" + \
               "我行代客境外理财产品正式推出" + \
               "2022" + \
               "2022-11" + \
               "我行首单ABS联席主承销商项目鹭盈2022年第一期个人消费贷款资产支持证券成功发行，为首家联主ABS项目的台资银行" + \
               "2022-09" + \
               "我行当选为中国银行业协会外资银行工作委员会第六届常务委员会副主任单位，陈峰行长当选外资银行工作委员会副主任" + \
               "2022-08" + \
               "我行制定ESG三年战略规划，推出“IMPACT”永续影响计划" + \
               "2022-08" + \
               "我行加入中国银行业协会私人银行与财富管理业务专业委员会并当选第三届常务委员会副主任单位，王文锋副行长当选副主任"
    cks = smart_chunk(question, task)
    print("智能切片结果为：", cks)
